//
// SoapTrace.Net - A soap message monitoring utility. 
//
// Copyright (C) 2005 Matthew Ward
//
// This program is free software; you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation; either version 2 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program; if not, write to the Free Software
// Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
//
// Matthew Ward (mrward@users.sourceforge.net)

using System;
using System.Collections;
using System.Net;
using System.Net.Sockets;
using WinPcap;

namespace SoapTrace.Core.Communication
{
	/// <summary>
	/// Represents the connection state.
	/// </summary>
	public enum TcpState
	{
		/// <summary>
		/// No state.
		/// </summary>
		None = 0,
		
		/// <summary>
		/// Syn data packet sent to initiate connection.
		/// </summary>
		SynSent = 1,
		
		/// <summary>
		/// Syn data packet acknowledged.  
		/// </summary>
		/// <remarks>
		/// Not used since Windows Raw Sockets send packets out of order
		/// and we cannot rely on the SYN, SYN-ACK, ACK sequence.
		/// </remarks>
		SynAcknowledged = 2,
		
		/// <summary>
		/// Syn received from server, connection has been made.
		/// </summary>
		SynReceived = 3
	}
	
	/// <summary>
	/// Represents a Tcp Connection.
	/// </summary>
	/// <remarks>
	/// It represents a connection made by the data packets involved
	/// and not at the higher level of abstraction with sockets.  For a higher
	/// level abstraction use the <see cref="TcpSocketConnection"/> class.
	/// </remarks>
	public class TcpConnection
	{		
		/// <summary>
		/// Bytes transferred event.
		/// </summary>
		public event TcpCommunicationEventHandler BytesTransferred;
		
		/// <summary>
		/// Connection event.
		/// </summary>
		public event TcpCommunicationEventHandler Connected;

		/// <summary>
		/// Disconnection event.
		/// </summary>
		public event TcpCommunicationEventHandler Disconnected;
		
		/// <summary>
		/// The server's IP address.
		/// </summary>
		private IPEndPoint m_ServerIPAddress;
		
		/// <summary>
		/// The next server's sequence number.
		/// </summary>
		private uint m_ServerSequenceNumber;
		
		/// <summary>
		/// The client's IP address.
		/// </summary>
		private IPEndPoint m_ClientIPAddress;
		
		/// <summary>
		/// The next client's sequence number.
		/// </summary>
		private uint m_ClientSequenceNumber;
		
		/// <summary>
		/// The current connection state.
		/// </summary>
		private TcpState m_TcpState = TcpState.None;
		
		/// <summary>
		/// Cached server tcp packets when packets have been lost.
		/// </summary>
		private ArrayList m_CachedServerPackets = new ArrayList( );
		
		/// <summary>
		/// Cached client tcp packets when packets have been lost.
		/// </summary>
		private ArrayList m_CachedClientPackets = new ArrayList( );
		
		/// <summary>
		/// The range of sequence numbers that will be cached.
		/// </summary>
		private uint m_CachedSequenceNumberRange = 65535;		
		
		/// <summary>
		/// Creates a new instance of the <see cref="TcpConnection"/>
		/// class.
		/// </summary>
		/// <param name="client">The client IP address.</param>
		/// <param name="clientSequenceNumber">The client's next sequence
		/// number.</param>
		/// <param name="server">The server IP address.</param>
		/// <param name="serverSequenceNumber">The server's next sequence 
		/// number.</param>
		[ CLSCompliant( false ) ] 
		public TcpConnection( IPEndPoint client, uint clientSequenceNumber, IPEndPoint server, uint serverSequenceNumber )
		{
			ServerIPAddress = server;
			ClientIPAddress = client;
			ServerSequenceNumber = serverSequenceNumber;
			ClientSequenceNumber = clientSequenceNumber;
			State = TcpState.None;
		}
		
		/// <summary>
		/// The IP address of the server.
		/// </summary>
		public IPEndPoint ServerIPAddress
		{
			get
			{
				return m_ServerIPAddress;
			}
			
			set
			{
				m_ServerIPAddress = value;
			}
		}
		
		/// <summary>
		/// Next server sequence number.
		/// </summary>
		[ CLSCompliant( false ) ] 
		public uint ServerSequenceNumber
		{
			get
			{
				return m_ServerSequenceNumber;
			}
			
			set
			{
				m_ServerSequenceNumber = value;
			}
		}
		
		/// <summary>
		/// The IP address of the client.
		/// </summary>
		public IPEndPoint ClientIPAddress
		{
			get
			{
				return m_ClientIPAddress;
			}
			
			set
			{
				m_ClientIPAddress = value;
			}
		}
		
		/// <summary>
		/// Next client sequence number.
		/// </summary>
		[ CLSCompliant( false ) ] 
		public uint ClientSequenceNumber
		{
			get
			{
				return m_ClientSequenceNumber;
			}
			
			set
			{
				m_ClientSequenceNumber = value;
			}
		}
		
		/// <summary>
		/// The current state of the connection.
		/// </summary>
		public TcpState State
		{
			get
			{
				return m_TcpState;
			}
			
			set
			{
				m_TcpState = value;
			}
		}
		
		/// <summary>
		/// Determines whether the <paramref name="packet"/> is a syn
		/// packet.
		/// </summary>
		/// <param name="packet">A <see cref="TcpPacket"/>.</param>
		/// <returns><see langword="true"/> if the packet is a Syn; 
		/// otherwise <see langword="false"/>.</returns>
		public static bool IsSynPacket( TcpPacket packet )
		{
			bool IsSyn = false;
			
			TcpControlTypes ControlFlag = ( TcpControlTypes )packet.ControlFlag;
				
			if( ControlFlag == TcpControlTypes.Syn )
			{
				IsSyn = true;
			}
		
			return IsSyn;
		}
		
		/// <summary>
		/// Determines whether the packet is a Syn acknowledgement.
		/// </summary>
		/// <param name="packet">A tcp packet.</param>
		/// <returns><see langword="true"/> if the packet is a Syn 
		/// acknowledgement; otherwise <see langword="false"/>.</returns>
		public static bool IsSynAckPacket( TcpPacket packet )
		{
			bool IsSynAck = false;
			
			TcpControlTypes ControlFlag = ( TcpControlTypes )packet.ControlFlag;
				
			if( ControlFlag == ( TcpControlTypes.Syn | TcpControlTypes.Ack ) )
			{
				IsSynAck = true;
			}
			
			return IsSynAck;
		}
		
		/// <summary>
		/// Determines whether the packet is an acknowledgement.
		/// </summary>
		/// <param name="packet">A tcp packet.</param>
		/// <returns><see langword="true"/> if the packet is a Ack 
		/// acknowledgement; otherwise <see langword="false"/>.</returns>
		public static bool IsAckPacket( TcpPacket packet )
		{
			bool IsAck = false;
			
			TcpControlTypes ControlFlag = ( TcpControlTypes )packet.ControlFlag;
				
			if( ControlFlag == TcpControlTypes.Ack )
			{
				IsAck = true;
			}
			
			return IsAck;
		}	
		
		/// <summary>
		/// Determines whether the packet is a Rst acknowledgement.
		/// </summary>
		/// <param name="packet">A tcp packet.</param>
		/// <returns><see langword="true"/> if the packet is a Rst 
		/// acknowledgement; otherwise <see langword="false"/>.</returns>
		public static bool IsRstAckPacket( TcpPacket packet )
		{
			bool IsSynAck = false;
			
			TcpControlTypes ControlFlag = ( TcpControlTypes )packet.ControlFlag;
				
			if( ControlFlag == ( TcpControlTypes.Rst | TcpControlTypes.Ack ) )
			{
				IsSynAck = true;
			}
			
			return IsSynAck;
		}		
		
		/// <summary>
		/// Determines whether the packet is an Fin.
		/// </summary>
		/// <param name="packet">A tcp packet.</param>
		/// <returns><see langword="true"/> if the packet is the Fin bit
		/// is set; otherwise <see langword="false"/>.</returns>
		public static bool IsFinPacket( TcpPacket packet )
		{
			bool IsAck = false;
			
			TcpControlTypes ControlFlag = ( TcpControlTypes )packet.ControlFlag;
				
			if( ( ControlFlag & TcpControlTypes.Fin ) > 0 )
			{
				IsAck = true;
			}
			
			return IsAck;
		}	
		
		/// <summary>
		/// Determines whether the packet is a Psh acknowledgement.
		/// </summary>
		/// <param name="packet">A tcp packet.</param>
		/// <returns><see langword="true"/> if the packet is a Psh 
		/// acknowledgement; otherwise <see langword="false"/>.</returns>
		public static bool IsPshAckPacket( TcpPacket packet )
		{
			bool IsSynAck = false;
			
			TcpControlTypes ControlFlag = ( TcpControlTypes )packet.ControlFlag;
				
			if( ControlFlag == ( TcpControlTypes.Psh | TcpControlTypes.Ack ) )
			{
				IsSynAck = true;
			}
			
			return IsSynAck;
		}
		
		/// <summary>
		/// Tests whether the <paramref name="testSequenceNumber"/>
		/// is in the range specified.
		/// </summary>
		/// <param name="testSequenceNumber">The sequence number to test.</param>
		/// <param name="rangeStartSequenceNumber">The sequence number
		/// at the start of the range.</param>
		/// <param name="range">The sequence number range.
		/// </param>
		/// <returns><see langword="true"/> if the number is in the 
		/// range; otherwise <see langword="false"/>.</returns>
		[ CLSCompliant( false ) ]
		public static bool IsSequenceNumberInRange( uint testSequenceNumber, uint rangeStartSequenceNumber, uint range )
		{
			bool IsInRange = false;
			
			if( IsRolloverRange( rangeStartSequenceNumber, range ) )
			{
				if( ( testSequenceNumber >= rangeStartSequenceNumber ) ||
				   ( testSequenceNumber < ( rangeStartSequenceNumber + range ) ) )
				{
					IsInRange = true;
				}
			}
			else
			{
				if( ( testSequenceNumber >= rangeStartSequenceNumber ) &&
				   ( testSequenceNumber < ( rangeStartSequenceNumber + range ) ) )
				{
				   	IsInRange = true;
				}
			}
			
			return IsInRange;
		}
		
		/// <summary>
		/// Determines whether the packet belongs to this connection.
		/// </summary>
		/// <param name="packet">A tcp packet.</param>
		/// <returns><see langword="true"/> if the packet belongs
		/// to this connection; otherwise <see langword="false"/>.</returns>
		public bool IsMatchingPacket( TcpPacket packet )
		{
			bool IsMatch = false;
			
			if( m_ClientIPAddress.Equals( packet.SourceAddress ) )
			{
				if( m_ServerIPAddress.Equals( packet.DestinationAddress ) )
				{
					IsMatch = true;
				}					
			}
			else if( m_ClientIPAddress.Equals( packet.DestinationAddress ) )
			{
				if( m_ServerIPAddress.Equals( packet.SourceAddress ) )
				{
					IsMatch = true;
				}
			}
							
			return IsMatch;
		}
		
		/// <summary>
		/// Extracts the data bytes from the packet and raises 
		/// the BytesTransferred event.
		/// </summary>
		/// <param name="packet">A tcp Psh/Ack packet.</param>
		/// <remarks>
		/// Windows raw sockets cause some complications.  The SYN,
		/// SYN-ACK, ACK connection packets are read at the client
		/// in the wrong order.  We see SYN, ACK, SYN-ACK.  So
		/// we have to allow this by only looking for the SYN-ACK
		/// from the remote server to indicate the 
		/// connection has been made.</remarks>
		public void ProcessPacket( TcpPacket packet )
		{
			switch( m_TcpState )
			{
				case TcpState.None:
				{
					if( TcpConnection.IsSynPacket( packet ) )
					{
						m_TcpState = TcpState.SynSent;	
					}
				}
				break;
				
				case TcpState.SynSent:
				{
					if( TcpConnection.IsSynAckPacket( packet ) )
					{
						SetClientConnectedState( packet.AcknowledgementNumber, 
						                        packet.DestinationAddress, 
						                        packet.SequenceNumber + 1, 
						                        packet.SourceAddress );
						
						OnConnected( packet.DestinationAddress, packet.SourceAddress );
					}
				}
				break;
				
				case TcpState.SynReceived:
				{
					if( TcpConnection.IsRstAckPacket( packet ) || 
					    TcpConnection.IsFinPacket( packet ) )
					{
						m_TcpState = TcpState.None;
						OnDisconnected( packet.SourceAddress, packet.DestinationAddress );
					}
					else if( TcpConnection.IsPshAckPacket( packet ) ||
					         TcpConnection.IsAckPacket( packet ) )
					{			
						ProcessDataPacket( packet );
					}							
				}
				break;				
			}
		}	
		
		/// <summary>
		/// Raises the <see cref="BytesTransferred"/> event.
		/// </summary>
		protected void OnBytesTransferred( IPEndPoint sourceAddress, IPEndPoint destinationAddress, byte[ ] bytes )
		{
			if( BytesTransferred != null )
			{
				TcpCommunicationEventArgs CommEventArgs = 
					new TcpCommunicationEventArgs( sourceAddress, destinationAddress, bytes );
				
				BytesTransferred( this, CommEventArgs );
			}
		}	
		
		/// <summary>
		/// Raises the <see cref="Connected"/> event.
		/// </summary>
		protected void OnConnected( IPEndPoint sourceAddress, IPEndPoint destinationAddress )
		{
			if( Connected != null )
			{
				TcpCommunicationEventArgs CommEventArgs = 
					new TcpCommunicationEventArgs( sourceAddress, destinationAddress );
				
				Connected( this, CommEventArgs );
			}
		}
		
		/// <summary>
		/// Raises the <see cref="Disconnected"/> event.
		/// </summary>
		protected void OnDisconnected( IPEndPoint sourceAddress, IPEndPoint destinationAddress )
		{
			if( Disconnected != null )
			{
				TcpCommunicationEventArgs CommEventArgs = 
					new TcpCommunicationEventArgs( sourceAddress, destinationAddress );
				
				Disconnected( this, CommEventArgs );
			}
		}	
		
		/// <summary>
		/// Extracts the data bytes from the packet and raises 
		/// the BytesTransferred event.
		/// </summary>
		/// <param name="packet">A tcp Psh/Ack packet.</param>
		private void ProcessDataPacket( TcpPacket packet )
		{
			if( packet.GetDataBytes( ).Length > 0 )
			{
				if( packet.SourceAddress.Equals( ClientIPAddress ) )
				{				
					ProcessDataPacket( packet, m_CachedClientPackets, ref m_ClientSequenceNumber );
				}
				else if( packet.SourceAddress.Equals( ServerIPAddress ) )
				{
					ProcessDataPacket( packet, m_CachedServerPackets, ref m_ServerSequenceNumber );
				}
			}
		}	
		
		/// <summary>
		/// Processes a data packet.
		/// </summary>
		/// <param name="packet">A tcp packet.</param>
		/// <param name="cachedPackets">The cached packets associated with this
		/// connection.</param>
		/// <param name="sequenceNumber">The current sequence number for the
		/// packets.
		/// </param>
		private void ProcessDataPacket( TcpPacket packet, ArrayList cachedPackets, ref uint sequenceNumber )
		{
			if( sequenceNumber == packet.SequenceNumber )
			{								
				int PacketLength = packet.GetDataBytes( ).Length;
				sequenceNumber += ( uint )PacketLength;
				
				OnBytesTransferred( packet.SourceAddress, packet.DestinationAddress, packet.GetDataBytes( ) );							
				
				// Check any cached packets.
				
				if( cachedPackets.Count > 0 )
				{
					ProcessCachedPackets( cachedPackets, ref sequenceNumber );
				}
			}
			else if( IsInRange( packet.SequenceNumber, sequenceNumber ) )
			{
				// Missed a tcp packet.
				
				cachedPackets.Add( packet );
			}
			else if( IsMergedDataPacket( sequenceNumber, packet.SequenceNumber, ( uint )( packet.GetDataBytes( ).Length ) ) )
			{
				System.Diagnostics.Debug.WriteLine( "Merged data packet!!!!" );
				ProcessMergedDataPacket( packet, cachedPackets, ref sequenceNumber );
			}
		}
		
		/// <summary>
		/// Updates the state after a client has connected.
		/// </summary>
		/// <param name="clientSequenceNumber">The client sequence number.</param>
		/// <param name="clientAddress">The client IP address.</param>
		/// <param name="serverSequenceNumber">The server sequence number.</param>
		/// <param name="serverAddress">The server IP address.</param>
		private void SetClientConnectedState( uint clientSequenceNumber, IPEndPoint clientAddress, uint serverSequenceNumber, IPEndPoint serverAddress )
		{
			m_TcpState = TcpState.SynReceived;
			m_ClientSequenceNumber = clientSequenceNumber;
			m_ClientIPAddress = clientAddress;
							
			m_ServerIPAddress = serverAddress;
			m_ServerSequenceNumber = serverSequenceNumber;
		}
		
		/// <summary>
		/// Looks at the cached packets to see if we can now
		/// raise events for the packets.
		/// </summary>
		private void ProcessCachedPackets( ArrayList cachedPackets, ref uint sequenceNumber )
		{
			for( int i = 0; i < cachedPackets.Count; ++i )
			{
				TcpPacket Packet = ( TcpPacket )cachedPackets[ i ];
				
				if( sequenceNumber == Packet.SequenceNumber )
				{
					int PacketLength = Packet.GetDataBytes( ).Length;
					sequenceNumber += ( uint )PacketLength;
							
					OnBytesTransferred( Packet.SourceAddress, Packet.DestinationAddress, Packet.GetDataBytes( ) );
				}
			}
			
			// Remove obsolete packets.
			int CurrentIndex = cachedPackets.Count;
			
			for( int j = 0; j < cachedPackets.Count; ++j )
			{
				--CurrentIndex;
				TcpPacket packet = ( TcpPacket )cachedPackets[ CurrentIndex ];
				
				if( false == IsInRange( packet.SequenceNumber, sequenceNumber ) )
				{
					// Remove
					cachedPackets.Remove( packet );
				}
			}
		}
		
		/// <summary>
		/// Checks whether the given sequence number should be cached.
		/// </summary>
		/// <param name="testSequenceNumber">The sequence number to check.</param>
		/// <param name="currentSequenceNumber">The next expected sequence number.</param>
		private bool IsInRange( uint testSequenceNumber, uint currentSequenceNumber )
		{
			return IsSequenceNumberInRange( testSequenceNumber, 
			                                currentSequenceNumber, 
			                                m_CachedSequenceNumberRange );
		}
		
		/// <summary>
		/// Checks whether the addition of the sequence number and the
		/// range causes the uint to rollover.
		/// </summary>
		/// <param name="sequenceNumber">A packet's sequence number.</param>
		/// <param name="range">The range of allowed sequence numbers.</param>
		/// <returns><see langword="true"/> if the range causes a rollover;
		/// otherwise <see langword="false"/>.</returns>
		private static bool IsRolloverRange( uint sequenceNumber, uint range )
		{
			bool IsRollover = false;
			
			if( ( sequenceNumber + range ) < sequenceNumber )
			{
				IsRollover = true;
			}
			
			return IsRollover;
		}	
		
		/// <summary>
		/// Determines whether the data packet represents a merged data packet
		/// and we need the data.
		/// </summary>
		private static bool IsMergedDataPacket( uint nextSequenceNumber, uint packetSequenceNumber, uint packetLength )
		{
			return IsSequenceNumberInRange( nextSequenceNumber, 
				packetSequenceNumber, 
				packetLength );
		}
		
		/// <summary>
		/// Takes a data packet that we have received a small amount of before
		/// and extract the bytes we have not received and call
		/// OnBytesReceived with the data.
		/// </summary>
		/// <param name="packet">A tcp packet.</param>
		/// <param name="cachedPackets">The cached packets associated with this
		/// connection.</param>
		/// <param name="sequenceNumber">The current sequence number for the
		/// packets.
		/// </param>
		private void ProcessMergedDataPacket( TcpPacket packet, ArrayList cachedPackets, ref uint sequenceNumber )
		{
			// Extract bytes.
			uint PacketLength = ( uint )( packet.GetDataBytes( ).Length );
			uint BytesToExtract = PacketLength - ( sequenceNumber - packet.SequenceNumber );
			
			byte[ ] ExtractedBytes = new byte[ BytesToExtract ];
			
			Array.Copy( packet.GetDataBytes( ), 
			           PacketLength - BytesToExtract,
			           ExtractedBytes,
			           0,
			           BytesToExtract );
			
			sequenceNumber += ( uint )BytesToExtract;

			OnBytesTransferred( packet.SourceAddress, packet.DestinationAddress, ExtractedBytes );
			
			// Check any cached packets.
			
			if( cachedPackets.Count > 0 )
			{
				ProcessCachedPackets( cachedPackets, ref sequenceNumber );
			}			
		}

	}
}
